import asyncio
import time
import os
import re
from typing import List, Dict, Any
from bilibili_api import video, search
from bilibili_api.search import SearchObjectType
from bilibili.processing import filter_danmaku, save_to_excel, generate_word_cloud
from config import MAX_VIDEOS

# 创建数据目录
os.makedirs('data', exist_ok=True)




async def main():
    search_keyword = input('请输入要搜索的关键词（如"大模型"）: ').strip()
    if not search_keyword:
        print("错误：搜索关键词不能为空")
        return

    print(f"\n开始爬取关键词 '{search_keyword}' 的视频数据...")
    
    # 搜索视频
    page = 1
    all_videos = []
    video_infos = []
    all_danmakus = []
    
    while len(all_videos) < MAX_VIDEOS:
        try:
            # 搜索视频
            search_result = await search.search_by_type(
                search_keyword,
                search_type=SearchObjectType.VIDEO,
                page=page
            )
            
            if not search_result or 'result' not in search_result:
                break
                
            results = search_result['result']
            if not results:
                break
                
            # 获取视频信息和弹幕
            for item in results:
                if len(all_videos) >= MAX_VIDEOS:
                    break
                    
                try:
                    bvid = item.get('bvid')
                    if not bvid:
                        continue
                        
                    print(f"正在获取视频 {bvid} 的信息和弹幕...")
                    v = video.Video(bvid=bvid)
                    
                    # 获取视频信息
                    info = await v.get_info()
                    all_videos.append(info)
                    
                    # 获取视频弹幕
                    danmakus = await v.get_danmakus()
                    all_danmakus.extend([d.text for d in danmakus])
                    
                    # 添加延迟避免请求过快
                    await asyncio.sleep(1)
                    
                except Exception as e:
                    print(f"获取视频信息失败: {e}")
                    continue
            
            page += 1
            print(f"已获取 {len(all_videos)} 个视频")
            await asyncio.sleep(1)  # 页面间延迟
            
        except Exception as e:
            print(f"搜索失败: {e}")
            break
    
    if not all_videos:
        print("未找到任何视频")
        return
    
    print(f"\n共获取到 {len(all_videos)} 个视频的数据")
    print(f"共获取到 {len(all_danmakus)} 条原始弹幕")
    
    # 处理弹幕
    filtered_danmakus = filter_danmaku(all_danmakus)
    print(f"过滤后剩余 {len(filtered_danmakus)} 条有效弹幕")
    
    # 保存数据
    timestamp = time.strftime('%Y%m%d_%H%M%S')
    
    # 保存视频信息
    video_excel = f"data/{search_keyword}_videos_{timestamp}.xlsx"
    save_to_excel(all_videos, video_excel, 'videos')
    print(f"视频信息已保存到: {video_excel}")
    
    # 保存弹幕数据
    danmaku_excel = f"data/{search_keyword}_danmaku_{timestamp}.xlsx"
    danmaku_data = [{'danmaku': d} for d in filtered_danmakus]
    save_to_excel(danmaku_data, danmaku_excel, 'danmaku')
    print(f"弹幕数据已保存到: {danmaku_excel}")
    
    # 生成词云
    wordcloud_file = f"data/{search_keyword}_wordcloud_{timestamp}.png"
    generate_word_cloud(filtered_danmakus, wordcloud_file)
    print(f"词云图已保存到: {wordcloud_file}")


if __name__ == "__main__":
    asyncio.run(main())